import unittest
from alg import NumSquares
from gen import TestGenerator


class TestNumSquares(unittest.TestCase):
    def setUp(self):
        self.solver = NumSquares()
        self.test_cases = TestGenerator.generate_test_cases()

        # 预处理已知测试用例的expected值
        known_cases = {
            12: 3,
            13: 2,
            1: 1,
            2: 2,
            3: 3,
            4: 1,
            5: 2
        }

        for i, (n, _) in enumerate(self.test_cases):
            if n in known_cases:
                self.test_cases[i] = (n, known_cases[n])

    def test_num_squares(self):
        for n, expected in self.test_cases:
            if expected is not None:
                with self.subTest(n=n):
                    result = self.solver.solve(n)
                    self.assertEqual(result, expected, f"Failed for n={n}")
            else:
                # 对于随机生成的测试用例，可以打印结果但不失败
                result = self.solver.solve(n)
                print(f"Tested n={n}, result={result}")


if __name__ == '__main__':
    unittest.main()